package com.keyue.isolate.config;


import com.keyue.isolate.context.DataIsolateContextHolder;
import com.keyue.isolate.context.DataIsolateValue;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * 针对线程池进行增强
 * 传递数据隔离上下文
 */
@Component
public class ThreadPoolDataIsolateContextConfig implements InitializingBean, ApplicationContextAware {

    private ApplicationContext applicationContext;

    @Override
    public void afterPropertiesSet() throws Exception {
        Map<String, ThreadPoolTaskExecutor> type2Bean = applicationContext.getBeansOfType(ThreadPoolTaskExecutor.class);
        for (Map.Entry<String, ThreadPoolTaskExecutor> entry : type2Bean.entrySet()) {
            //先硬编码 后续考虑配置化
            if (entry.getKey().equals("exportExecutor")) {
                entry.getValue().setTaskDecorator(new ContextCopyDecorator());
            }
        }
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }


    /**
     * 线程池任务-上下文拷贝 装饰器
     */
    static class ContextCopyDecorator implements TaskDecorator {


        @Override
        public Runnable decorate(Runnable runnable) {
            if (!DataIsolateContextHolder.hasContext()) {
                return runnable;
            }
            List<DataIsolateValue> dataIsolateValues = DataIsolateContextHolder.get();

            return () -> {
                try {
                    DataIsolateContextHolder.setContextValue(dataIsolateValues);
                    runnable.run();
                } finally {
                    DataIsolateContextHolder.clear();
                }
            };
        }
    }
}
